# Context Engineering System：程序主控系统 
from typing import Dict
import re
# from apps.cpm.tes.snr_range_agent import SnrRangeAgent
from apps.cpm.rlm_context import RlmContext
from apps.cpm.agent.snr_ranges_agent import SnrRangesAgent
from apps.cpm.agent.tau_snr_agent import TauSnrAgent
from apps.cpm.agent.ref_snr_agent import RefSnrAgent
from apps.cpm.agent.sr_azi_bw_agent import SrAziBwAgent
from apps.cpm.agent.pa_snr_agent import PaSnrAgent
from apps.cpm.agent.radar_design_v1_agent import RadarDesignV1Agent
from apps.cpm.mcp.mcpc_fm import McpcFm
from apps.cpm.cpm_engine import CpmEngine

class CeSystem(object):
    TESS = [SnrRangesAgent, TauSnrAgent, RefSnrAgent, SrAziBwAgent, PaSnrAgent, RadarDesignV1Agent]
    def __init__(self):
        self.name = 'apps.cpm.ce_system.CeSystem'

    def startup(self, user_query:str) -> Dict:
        context = RlmContext()
        context.current_input = user_query
        # 识别当前任务：确定是哪个Tes
        atom_prompt = self.generate_intent_prompt(user_query=context.current_input)
        print(f'{atom_prompt}\n\n')
        resp = CpmEngine.infer(query=atom_prompt)
        print(f'resp: {resp};\n\n')
        pattern = r'<\|intent_start\|>(.*?)<\|intent_end\|>'
        rsts = re.findall(pattern, resp, flags=re.DOTALL)
        ce_idx = int(rsts[-1])
        # arrs0 = resp.split('<|intent_start|>')
        # ce_idx = int(arrs0[1].split('<|intent_end|>')[0])
        ce_idxs = [ce_idx] # 有多个步骤，上步输入作为下步输出
        data = CeSystem.TESS[ce_idxs[0]].run(context=context)
        for idx in range(1, len(ce_idxs)):
            data = CeSystem.TESS[idx].run(data)
        return {
            'status': 'success',
            'report': './work/reports/v001/report.md',
            'result': data
        }
    
    def generate_intent_prompt(self, user_query:str) -> str:
        '''
        生成意图识别提示词
        '''
        prompt = '''
# 角色
假设你是资深雷达系统设计师

# 任务
你需要根据用户的需求描述，识别出用户的意图，即设计任务。意图列表为：
0: 信噪比与距离关系,
1: 脉冲宽度与信噪比关系
2: 由参考值求当前信噪比
3: 搜索雷达确定水平波束宽度
4: 脉冲累积下信噪比与距离
5: 雷达设计v1版

# 要求
输出仅包括：<|intent_start|>x<|intent_end|>，且这个序列仅输出一次,x=0,1...。

# 输出格式
<|intent_start|>1<|intent_end|>

# 示例
## 示例1
输入：假设一个工作频率为$5.6\\times 10^{9}\\text{Hz}$，天线增益为45.0dB，有效噪声温度为290.0K，带宽为$5.0\\times 10^{6}\\text{Hz}$，噪声系数为3.0dB，系统损耗为6.0dB，计算25Km到165Km范围内，分为1000个距离单元，对每个距离单元分别计算：1. 雷达截面积rCS分别为：$0.1\\text{m}^{2}$、$0.01\\text{m}^{2}$、$1.0\\text{m}^{2}$时，信噪比与距离之间的关系并绘制曲线；2. 雷达的峰值功率分别为：$1.5\\times 10^{6} \\text{W}$、$0.6\\times 10^{6}\\text{W}$、$2.7\\times 10^{6}\\text{W}$时，信噪比与距离之间的关系并绘制曲线；
输出：<|intent_start|>0<|intent_end|>
## 示例2
输入：假设一个工作频率为$5.6\times 10^{9}\\text{Hz}$，发射功率为$1.0\\times 10^{6}\\text{W} $，天线增益为40.0dB，雷达截面积为$0.1\\text{m}^{2}$，有效噪声温度为300.0K，噪声系数为5.0dB，系统损耗为6.0dB，距离范围为：75Km、100Km、150Km，研究信噪比范围为5dB到20dB，将其分为200等份，求出对应的最小脉冲宽度并绘制曲线。
输出：<|intent_start|>1<|intent_end|>
## 示例3
输入：某雷达参考距离为86Km，参考脉宽为0.1us，参考信噪比为20.0dB，参考雷达截面积为0.1$m^{2}$，处理损耗为2.0dB，当前距离为120Km，当前RCS为0.2$m^{2}$，当前脉冲宽度为0.1us，计算当前的信噪比。
输出：<|intent_start|>2<|intent_end|>
## 示例4
输入：设计一个地基雷达，可以识别10公里外的飞机和2公里外的导弹，最大探测距离为60公里。假设飞机的平均RCS为6dBsm，导弹平均RCS为-10dBsm。雷达水平搜索范围为$360^{\\circ}$，竖直搜索范围为$10^{\\circ}$。扫描率为2秒。距离精度为150米。假设噪声数为8dB，接收噪声为10dB。使用fan beam，水平波束宽度小于3度，SNR为15dB。请确定水平波束宽度。
输出：<|intent_start|>3<|intent_end|>
## 示例5
输入：一个MMW雷达规范如下：中心频率$f=94GHz$，脉冲宽度$\\tau=50\\times 10^{-9} s$，峰值功率$P_{t}=4W$，水平扫描范围$\\Theta_{A}=\\pm 120^{\\circ}$，脉冲重复频率$PRF=10KHz$，噪声数$F=7dB$，天线直径为$D=12in$，天线增益$G=47dB$，目标RCS为$\\sigma=20m^{2}$，系统损失$L=10dB$，雷达扫描时间$T_{sc}=3s$，温度设为$T_{e}=290K$。距离范围为从1公里到12公里并分为1000份。请计算：波长$\\lambda$；距离分辨率$\\bigtriangleup R$；带宽$B$；天线半功率波束宽度；天线扫描率；计算SNR为10dB时的探测距离；绘制距离与SNR关系的曲线；计算目标得到的脉冲数，当使用脉冲累积时的探测距离，假设SNR不变（与单脉冲时相同）。
输出：<|intent_start|>4<|intent_end|>
## 示例6
输入：设计一款满足如下功能的搜索和检测地基雷达。威胁包括飞机和导弹两大类。其中飞机的RCS为6dBsm（$\\sigma_{a}=4m^{2}$），导弹平均RCS为-3dBsm（$\\sigma_{m}=0.5m^{2}$）。导弹的高度是2Km，飞机的高度为7Km。水平扫描范围为360度，扫描速度为2秒钟转一圈。采用L到X波段，距离分辨率为150m。角度分辨率暂时不要求。危胁由一个飞机和一个导弹构成。噪声数为$F=6dB$，总接收损失$L=8dB$。采用fan beam方式，水平扫描波束宽度小于3度。假设13dB的信噪比是检测阈值，净空距离为$R_{min}=30Km$，防空导弹的速度为$v_{i}=250m/s$，飞机的速度为$v_{a}=400m/s$，导弹速度为$v_{m}=150m/s$，并且地面是平的。
输出：<|intent_start|>5<|intent_end|>

# 当前输入
'''
        prompt = f'{prompt}\n{user_query}'
        return prompt